#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <bitset>
#include <vector>
#include <map>
#include <cmath>
#include "dct.h"
#include "common.h"
#include "decode.h"

using namespace std;

vector<bool> read_jpg(std::string filename, int &fileSize)
{
    streampos size;
    vector<bool> rawData;
    char * memblock;
    
    ifstream file(filename.c_str(),ios::in|ios::binary|ios::ate);
    if (file.is_open())
    {
        size = file.tellg();
        fileSize = size * 8;
        //cout << "the file size is: " << size << endl;
        
        file.seekg (0, ios::beg);
        for(int i=0; i<size; i++){
            memblock = new char;
            *memblock = file.get();
            //cout << *memblock <<endl;
            bitset<8> byte = bitset<8>(* memblock);
            //cout << byte << endl;
            for (int i=0; i<8; i++){
                //cout << byte[7-i] << endl;
                rawData.push_back(byte[7-i]);
            }
            //cout << endl;
            //file.seekg(1,ios::cur);
            delete memblock;
        }
        file.close();
        
        //cout << "the entire file content is in memory" << endl;
        
    }
    else cout << "Unable to open file";
    return rawData;
}

void decoder_codebook(map<vector<bool>, int> &dc_codebook, map<vector<bool>, pair<int,int> > &ac_codebook)
{
    ifstream codebookfile(CODEBOOK);
    string line;
    int count = 1;
    vector<bool> code;
    int run;
    int cat;
    
    while (getline(codebookfile, line)) {
        if (line == "") continue;
        else if (count <= 12){
            cat = charToInt(line[0]);
            for (int i=1; i!=line.length(); ++i)
            {
                if(line[i]=='0') code.push_back(0);
                if(line[i]=='1') code.push_back(1);
            }
            dc_codebook.insert(pair<vector<bool>, int>(code, cat));
        }
        else
        {
            run = charToInt(line[0]);
            cat = charToInt(line[2]);
            for (int i=3; i!=line.length(); ++i)
            {
                if(line[i]=='0') code.push_back(0);
                if(line[i]=='1') code.push_back(1);
            }
            ac_codebook.insert(pair<vector<bool>, pair<int, int> > (code, pair<int,int>(run, cat)));
        }
        code.clear();
        count++;
    }
    codebookfile.close();
}

int binaryToBase10(int n)
{
    int output = 0;
    
    for(int i=0; n > 0; i++) {
        
        if(n % 10 == 1) {
            output += (1 << i);
        }
        n /= 10;
    }
    
    return output;
}

void huffmanDC(map<vector<bool>, int> &dc_codebook, vector<bool> &rawData, vector<int> &intBuffer, int &index, int &previousDC)
{
    vector<bool> code;
    while (true) {
        //find the code in codebook
        code.push_back(rawData.at(index));
        //cout<<rawData.at(index);
        //if has such category, get the delta value and compute the first value in the block
        if (dc_codebook.find(code)!=dc_codebook.end()) {
            //cout << endl;
            //cout << "DC code: ";
            //for (int i =0; i < code.size(); ++i) {
            //    cout << code[i];
            //}
            
            int cat=dc_codebook[code];
            //cout << " Category: " << cat << endl;
            int deltaDCBit = 0;
            int maxDelta = 0;
            //cout << "DC coefficient raw Data: ";
            for (int i = 0; i<cat; i++) {
                index ++;
                deltaDCBit = deltaDCBit * 10 + rawData.at(index);
                maxDelta  += (1 << i);
                //cout << rawData[index];
            }
            
            //find delta range
            int deltaDC = binaryToBase10(deltaDCBit);
            if (deltaDC < ((maxDelta+1)/2)) {
                deltaDC -= maxDelta;
                //cout << "coefficent range after minus: " << deltaDC << endl;
            }
            int currentDC = previousDC + deltaDC;
            previousDC = currentDC;
            
            //cout << " coefficient range: " << currentDC << endl;
            intBuffer.push_back(currentDC);
            index++;
            break;
        }
        index++;
    }
}


void huffmanAC(map<vector<bool>, pair<int, int> > &ac_codebook, vector<bool> &rawData, vector<int> &intBuffer, int &index)
{
    //find the code in codebook
    vector<bool> code;
    while (intBuffer.size() < 64) {
        code.push_back(rawData.at(index));
        index++;
        if (ac_codebook.find(code)!=ac_codebook.end()) {
            //cout << "AC code: ";
            //for (int i =0; i < code.size(); ++i) {
            //    cout << code[i];
            //}
            pair<int, int> runCat = ac_codebook[code];
            int run = get<0>(runCat);
            int cat = get<1>(runCat);
            //cout << " run: " << run << " cat: " << cat << endl;
            if (run == 0 && cat == 0) {
                //fill in zero to 64bit
                int noZero = 64 - intBuffer.size();
                for (int i = 0; i < noZero; i++) {
                    intBuffer.push_back(0);
                }
            } else if (run == 15 && cat == 0) {
                // if F, 0, add 15 zero
                for (int i = 0; i < 15; i++) {
                    intBuffer.push_back(0);
                }
            } else {
                // add run no. of zero, then add category
                for (int i=0; i<run; i++) {
                    intBuffer.push_back(0);
                }
                int ACValueInBit = 0;
                int maxACValue = 0;
                //cout << "raw data for AC cat value: ";
                for (int i=0; i<cat; i++) {
                    ACValueInBit = ACValueInBit * 10 + rawData.at(index);
                    maxACValue  += (1 << i);
                    //cout << rawData[index];
                    index++;
                }
                int ACValue = binaryToBase10(ACValueInBit);
                if (ACValue < ((maxACValue+1)/2)) {
                    ACValue -= maxACValue;
                }
                //cout << " AC value: " << ACValue << endl;
                intBuffer.push_back(ACValue);
            }
            code.clear();
        }
    }
    
    //for (int i=0;i!=intBuffer.size();++i){
    //    cout<<intBuffer[i];
    //}
    //cout<<endl;
    if (intBuffer.size() > 64) {
        cout << "ERROR: WHAT HAPPENED????" << endl;
    }
}

int **scan(vector<int> &intBuffer)
{
    int scanM[BLK_DIM][BLK_DIM]=
    {   {0,  1,  5,  6,  14, 15, 27, 28},
        {2,  4,  7,  13, 16, 26, 29, 42},
        {3,  8,  12, 17, 25, 30, 41, 43},
        {9,  11, 18, 24, 31, 40, 44, 53},
        {10, 19, 23, 32, 39, 45, 52, 54},
        {20, 22, 33, 38, 46, 51, 55, 60},
        {21, 34, 37, 47, 50, 56, 59, 61},
        {35, 36, 48, 49, 57, 58, 62, 63}};
    
    int **matrix8x8 = new int*[BLK_DIM];
    
    for (int i=0;i!=BLK_DIM;++i){
        matrix8x8[i] = new int[BLK_DIM];
        for (int j=0;j!=BLK_DIM;++j)
            matrix8x8[i][j]=intBuffer[scanM[i][j]];
    }
    return matrix8x8;
    
}

float **deQuantize(int **matrix8x8, int quantize_matrix[BLK_DIM][BLK_DIM], float qs)
{
    float ** deQuantized_matrix = new float* [BLK_DIM];
    for (int i=0;i!=BLK_DIM;++i) deQuantized_matrix[i] = new float [BLK_DIM];
    for (int i = 0; i < BLK_DIM; ++i) {
        for (int j = 0; j < BLK_DIM; ++j) {
            deQuantized_matrix[i][j] = (float)(matrix8x8[i][j] * quantize_matrix[i][j] * qs);
        }
    }
    return deQuantized_matrix;
}

float *convertImg(vector<vector<float> > processedBuffer, int dim)
{
    float *img = new float[dim*dim];
    int rowBlk = dim/BLK_DIM;
    int colBlk = dim/BLK_DIM;
    int imgPos = 0;
    for (int i=0; i < rowBlk; ++i) {
        for (int blkRow=0; blkRow<BLK_DIM; ++blkRow) {
            for (int j=0; j<colBlk; ++j) {
                
                vector<float> blk = processedBuffer[i*colBlk+j];
                //for (int z = 0; z<blk.size(); ++z) {
                //    cout << blk[z] << " ";
                //}
                //cout <<endl;
                for (int blkCol=0; blkCol<BLK_DIM; ++blkCol) {
                    *(img+imgPos) = blk[blkRow*BLK_DIM+blkCol];
                    imgPos += 1;
                }
            }
        }
    }
    return img;
}

void writeFile(int dim, string outputFile, float *img)
{
    int imgPos = 0;
    ofstream fs(outputFile, fstream::out);
    fs << "P5\n";
    fs << "512 512\n";
    fs << "255\n";
    for (int i = 0; i<dim; ++i) {
        for (int j=0; j<dim; ++j) {
            //std::cout << ((unsigned long)(*(img+imgPos))) << std::endl;
            std::bitset<8> byte = std::bitset<8>((unsigned long)(*(img+imgPos)));
            fs.write(reinterpret_cast<char *>(&byte),1);
            imgPos += 1;
        }
    }
    fs.close();
}

int decode(float qs){
    //read jpg file into rawData
    vector<bool> rawData;
    int fileSize;
    int dim = 512;
    rawData = read_jpg(COMPRESSED_IMG, fileSize);
//    for (int i = 0; i < rawData.size(); ++i) {
//        cout << rawData[i];
//        if (i%7 == 0 and i != 0) {
//           cout << "\n" <;
//        }
//    }
//    cout << "\n" <<endl;
    
    //load the given code book
    map<vector<bool>, int> dc_codebook;
    map<vector<bool>, pair<int,int> > ac_codebook;
    decoder_codebook(dc_codebook,ac_codebook);
//    for (std::map<vector<bool>, int >::iterator it=dc_codebook.begin();it!=dc_codebook.end();++it)
//    {
//        
//        for (std::vector<bool>::iterator it2=(it->first).begin();it2!=(it->first).end();++it2)
//        {
//            std::cout<<*(it2);
//        }
//        std::cout<<" : " << it->second<< endl;
//    }
    
    int index = 0;
    vector<vector<float> > processedBuffer;
    int countBlk = 0;
    int previousDC = 0;
    while (countBlk < (dim/BLK_DIM)*(dim/BLK_DIM)) {
        countBlk += 1;
        //cout << "Blk No. : " << countBlk << endl;
        //cout << "current index: " << index << endl;
        vector<int> intBuffer;
        int **matrix8x8;
    
        //decode huffman DC, AC
        huffmanDC(dc_codebook, rawData, intBuffer, index, previousDC);
        //cout << "After DC: ";
        //for (int i = 0; i<intBuffer.size(); ++i) {
            //cout << intBuffer[i];
        //}
        //cout << endl;
        huffmanAC(ac_codebook, rawData, intBuffer, index);
        //cout << "After AC: ";
        //for (int i = 0; i<intBuffer.size(); ++i) {
            //cout << intBuffer[i];
        //}
        //cout << "\n" << endl;
        
        //scan
        matrix8x8 = scan(intBuffer);
    
        //de-quantize
        int quantize_matrix[BLK_DIM][BLK_DIM]={{16,11,10,16,24,40,51,61},
                                {12,12,14,19,26,58,60,55},
                                {14,13,16,24,40,57,69,59},
                                {14,17,22,29,51,87,80,62},
                                {18,22,37,56,68,109,103,77},
                                {24,35,55,64,81,104,113,92},
                                {49,64,78,87,103,121,120,101},
                                {72,92,95,98,112,110,103,99}};
        float **deQuantize8x8 = deQuantize(matrix8x8 , quantize_matrix, qs);
        vector<float> deQuantizedVector;
        for (int i=0;i!=BLK_DIM;++i){
            for (int j=0;j!=BLK_DIM;++j){
                deQuantizedVector.push_back(deQuantize8x8[i][j]);
            }
        }
    
        //temp
//        for (int i = 0; i < (dim/BLK_DIM)*(dim/BLK_DIM); ++i) {
//            processedBuffer.push_back(deQuantizedVector);
//        }
        processedBuffer.push_back(deQuantizedVector);
    
        for(int i = 0; i < BLK_DIM; ++i) {
            delete [] matrix8x8[i];
            delete [] deQuantize8x8[i];
        }
        delete [] matrix8x8;
        delete [] deQuantize8x8;
    }
    
    //convert into one image
    float *img = convertImg(processedBuffer, dim);
    
    //inverse DCT
    dct(img,dim,BLK_DIM,-1);

    //add 128
    for (int i=0; i < dim*dim; ++i) {
        *(img+i) += 128;
        //std::cout << *(img+i) << " ";
    }
    
    //write to file
    writeFile(dim, RECONSTRUCTED_IMG, img);
    delete [] img;
    return 0;
}
